#include "pch.h"
#include "ThreadpoolWait.h"
#include "AsyncLibImpl.h"

namespace AsyncLib
{

CThreadpoolWait::CThreadpoolWait()
{
	ZeroMemory(m_pTpCaches, sizeof(m_pTpCaches));
}

CThreadpoolWait::~CThreadpoolWait()
{
	assert(m_mapTasks.empty());
}

void CThreadpoolWait::Init()
{
	m_mapTasks.reserve(ARRAYSIZE(m_pTpCaches));
}

void CThreadpoolWait::Stop()
{
	std::unique_lock lock(m_mutex);

	for (auto& [pTpWait, _] : m_mapTasks)
	{
		SetThreadpoolWait(pTpWait, NULL, NULL);
	}
}

void CThreadpoolWait::Exit()
{
	// They have been waited and closed outside
	ZeroMemory(m_pTpCaches, sizeof(m_pTpCaches));
	m_mapTasks.clear();
}

void CThreadpoolWait::AddTask(HANDLE hHandle, UINT nMilliseconds, std::function<void(DWORD)>&& fnCallback)
{
	auto pTpWait = GetTpWait();
	assert(pTpWait);

	FILETIME ftTimeout;
	(INT64&)ftTimeout = -(nMilliseconds * 10000i64);

	std::unique_lock lock(m_mutex);
	m_mapTasks.emplace(pTpWait, std::move(fnCallback));
	SetThreadpoolWait(pTpWait, hHandle, nMilliseconds == INFINITE ? NULL : &ftTimeout);
}

void CThreadpoolWait::OnCallback(PTP_WAIT pTpWait, DWORD dwResult)
{
	g_asyncLib.SetInThreadPool();
	std::unique_lock lock(m_mutex);

	if (g_asyncLib.IsRunning())
	{
		ReleaseTpWait(pTpWait);
		auto node = m_mapTasks.extract(pTpWait);
		lock.unlock();
		assert(node);

		if (node)
			node.mapped()(dwResult);
	}
}

PTP_WAIT CThreadpoolWait::GetTpWait()
{
	for (int i = 0; i != ARRAYSIZE(m_pTpCaches); i++)
	{
		auto pTpWait = (PTP_WAIT)InterlockedExchangePointer((PVOID *)&m_pTpCaches[i], NULL);
		if (pTpWait)
			return pTpWait;
	}

	return CreateThreadpoolWait([](PTP_CALLBACK_INSTANCE, PVOID pContext, PTP_WAIT pTpWait, TP_WAIT_RESULT dwResult)
	{
		((CThreadpoolWait *)pContext)->OnCallback(pTpWait, dwResult);
	}, this, g_asyncLib.GetCallbackEnv());
}

void CThreadpoolWait::ReleaseTpWait(PTP_WAIT pTpWait)
{
	for (int i = 0; i != ARRAYSIZE(m_pTpCaches); i++)
	{
		pTpWait = (PTP_WAIT)InterlockedExchangePointer((PVOID *)&m_pTpCaches[i], pTpWait);
		if (pTpWait == NULL)
			return;
	}

	CloseThreadpoolWait(pTpWait);
}

} // namespace AsyncLib
